#include "BoundaryTools.hpp"
#include "Cell.hpp"
#include "MeshIO.hpp"
#include "Function.hpp"
#include "Quad.hpp"
#include "MakeHexMesh.hpp"
#include <Eigen/Core>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <algorithm>
#include <chrono>
#include <cmath>
#include <iostream>
#include <vector>

class RhsF : public Function<3> {
public:
  double operator() (const Eigen::Vector<double, 3> &point) const override{
    // return point.norm();
    return 1;
  }
};

class BoundaryF : public Function<3> {
public:
  double operator() (const Eigen::Vector<double, 3> &point) const override{
    // return cos(point[0])*cos(point[1])*cos(point[2]);
    return 0;
  }
};

class Poisson {
public:
  typedef Eigen::SparseMatrix<double> SpMat;
  typedef Eigen::VectorXd vector_type;

public:
  Poisson();
  void run();

private:
  void get_mesh();
  void setup_system();
  void solve();
  void output_results() const;

  Eigen::MatrixXd V; // `V`为顶点数\*3 double型矩阵，每行为一个顶点坐标
  Eigen::MatrixXi H; // `T`为四面体单元数\*4 整型矩阵，每行为四面体单元对应顶点编号

  SpMat system_matrix;
  vector_type solution;
  vector_type system_rhs;

  RhsF rhs_f;
  BoundaryF boundary_f;
};

Poisson::Poisson() {}

void Poisson::get_mesh() {
  readTetMeshFromVtk("./../tetmesh/sphere.vtk", V, H);
}

void Poisson::setup_system() {
  auto start_time = std::chrono::system_clock::now();

  int num_of_cells = H.rows();
  int num_of_nodes = V.rows();
  system_matrix.resize(num_of_nodes, num_of_nodes);
  system_rhs.resize(num_of_nodes);
  system_matrix.setZero();
  system_rhs.setZero();
  std::vector<Eigen::Triplet<double>> triplets_list;
  int fe_per_cell = 4;

  // 填充矩阵system_matrix
  for (int n = 0; n < num_of_cells; n++) {
    Eigen::Vector4i tetra = H.row(n); // 记录当前单元四面体的顶点标号
    TetCell<1> cell(V.row(tetra(0)),V.row(tetra(1)),V.row(tetra(2)),V.row(tetra(3)));
    TetQuad<3> quad_1;
    for (int i = 0; i < fe_per_cell; i++) {
      for (int j = 0; j < fe_per_cell; j++) {
        /* 进行数值积分计算 */
        double I=0;
        for (auto [w,q] : quad_1) {
          I += w*cell.jacobi()*cell.gradient(i, q).dot(cell.gradient(j, q));
        }
        triplets_list.push_back(
            Eigen::Triplet<double>(tetra(i), tetra(j),I));
      }
    }

    // 填充向量system_rhs
    for (int i = 0; i < fe_per_cell; ++i) {
      double I=0;
      for (auto [w,q] : quad_1) {
        I += w*cell.jacobi()*cell.value(i,q)*rhs_f(cell.map(q));
      }
      system_rhs(tetra(i)) += I;
    }
  }
  system_matrix.setFromTriplets(triplets_list.begin(), triplets_list.end());

  auto end_time = std::chrono::system_clock::now();
  std::cout << "set up basic system run time: " 
            << std::chrono::duration<double>{end_time-start_time}.count()
            << std::endl;

  // 设置边界条件
  start_time = std::chrono::system_clock::now();
  std::vector<int> boundary_nodes = TetTools::get_boundary(H);
  TetTools::set_boundary(boundary_f, boundary_nodes, V, system_matrix, system_rhs);
  end_time = std::chrono::system_clock::now();
  std::cout << "boundary condition run time: "
            << std::chrono::duration<double>{end_time - start_time}.count()
            << std::endl;
}

void Poisson::solve() {
  Eigen::ConjugateGradient<SpMat> solver;
  auto start_time = std::chrono::system_clock::now();
  solution = solver.compute(system_matrix).solve(system_rhs);
  auto end_time = std::chrono::system_clock::now();
  std::cout << "solver run time: "
            << std::chrono::duration<double>{end_time - start_time}.count()
            << std::endl;
  std::cout << "maximal element: "
            << *std::max_element(solution.begin(), solution.end()) << std::endl;
  std::cout << "minimal element: "
            << *std::min_element(solution.begin(), solution.end()) << std::endl;
}

void Poisson::output_results() const {
  writeTetMeshToVtk("solution_of_sphere.vtk", V, H,
                    {{AttribType::POINT, AttribFormat::SCALAR,
                      AttribDataType::DOUBLE, "solution", solution.data()}});

  std::cout << "nonzero entries of sys_mat: " << system_matrix.nonZeros()
            << std::endl;
}

void Poisson::run() {
  // make_hex_mesh(V, H, 32, 32, 32, Eigen::Vector3d{-1,-1,-1}, Eigen::Vector3d{1,1,1});
  get_mesh();
  setup_system();
  solve();
  output_results();
}

int main() {
  Poisson poisson_eqn;
  poisson_eqn.run();

  return 0;
}